package mongodb

import (
	"context"
	"fmt"
	"github.com/qiniu/qmgo"
	"go.mongodb.org/mongo-driver/bson"
)

const PH = 0

type Connection struct {
	*qmgo.Client
	ctx context.Context
	db  *qmgo.Database
}

var (
	defaultConnectTimeoutMS int64 = 30 * 1000
	defaultSocketTimeoutMS  int64 = 300 * 1000
)

func QuickConnect(uri string, dbName string) (*Connection, error) {
	return Connect(uri, dbName, defaultConnectTimeoutMS, defaultSocketTimeoutMS)
}

func Connect(uri string, dbName string, connectTimeoutMS int64, socketTimeoutMS int64) (*Connection, error) {
	if dbName == "" {
		return nil, fmt.Errorf("need database name")
	}

	ctx := context.Background()
	client, err := qmgo.NewClient(ctx, &qmgo.Config{
		Uri:              uri,
		Database:         dbName,
		ConnectTimeoutMS: &connectTimeoutMS,
		SocketTimeoutMS:  &socketTimeoutMS,
	})
	if err != nil {
		return nil, err
	}

	c := &Connection{
		Client: client,
		ctx:    ctx,
		db:     client.Database(dbName),
	}

	return c, nil
}

func (c *Connection) Context() context.Context {
	return c.ctx
}

func (c *Connection) DB() *qmgo.Database {
	return c.db
}

func (c *Connection) C(name string) *qmgo.Collection {
	return c.db.Collection(name)
}

func (c *Connection) Close() error {
	return c.Client.Close(c.ctx)
}

// EnsureCounter goroutine safe
func (c *Connection) EnsureCounter(collection string, id string) error {
	_, err := c.db.Collection(collection).InsertOne(c.ctx, bson.M{
		"_id": id,
		"seq": 0,
	})
	if !qmgo.IsDup(err) {
		return err
	}

	return nil
}

// NextSeq goroutine safe
func (c *Connection) NextSeq(collection string, id string) (int, error) {
	var res struct {
		Seq int
	}

	err := c.db.Collection(collection).Find(c.ctx, bson.D{{Key: "_id", Value: id}}).Apply(qmgo.Change{
		Update:    bson.M{"$inc": bson.M{"seq": 1}},
		ReturnNew: true,
	}, &res)

	return res.Seq, err
}

// EnsureIndexes goroutine safe
func (c *Connection) EnsureIndexes(collection string, uniques []string, indexes []string) error {
	return c.db.Collection(collection).EnsureIndexes(c.ctx, uniques, indexes)
}
